import pandas as pd
import pymysql
from dotenv import load_dotenv
import os
from datetime import datetime

# 加载环境变量
load_dotenv()

def get_db_connection():
    """获取数据库连接"""
    db_config = {
        'host': os.getenv('DB_HOST', 'localhost'),
        'port': int(os.getenv('DB_PORT', 3306)),
        'user': os.getenv('DB_USER', 'root'),
        'password': os.getenv('DB_PASSWORD', '123456'),
        'database': os.getenv('DB_NAME', 'zpsjfx'),
        'charset': 'utf8mb4'
    }
    return pymysql.connect(**db_config)

def create_table():
    """创建数据表"""
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            # 读取SQL文件
            with open(r'D:\\python1\\doc\\boss_jobs_schema.sql', 'r', encoding='utf-8') as f:
                sql_script = f.read()
            
            # 执行SQL脚本
            cursor.execute(sql_script)
            
            # 创建索引
            index_sqls = [
                "CREATE INDEX idx_title ON boss_jobs(title)",
                "CREATE INDEX idx_salary ON boss_jobs(salary)",
                "CREATE INDEX idx_company ON boss_jobs(companyTitle)",
                "CREATE INDEX idx_address ON boss_jobs(address)",
                "CREATE INDEX idx_type ON boss_jobs(type)",
                "CREATE INDEX idx_educational ON boss_jobs(educational)",
                "CREATE INDEX idx_work_experience ON boss_jobs(workExperience)"
            ]
            
            for index_sql in index_sqls:
                try:
                    cursor.execute(index_sql)
                except Exception as e:
                    print(f"创建索引时出错: {e}")
            
        connection.commit()
        print("数据表创建成功")
    except Exception as e:
        print(f"创建数据表时出错: {e}")
    finally:
        connection.close()

def import_csv_to_db():
    """将CSV数据导入数据库"""
    # 读取CSV文件
    file_path = r'D:\python1\doc\boss直聘数据.csv'
    try:
        df = pd.read_csv(file_path, encoding='utf-8')
        print(f"成功读取CSV文件，共{len(df)}条记录")
    except Exception as e:
        print(f"读取CSV文件失败: {e}")
        return
    
    # 连接数据库
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            # 清空现有数据
            cursor.execute("DELETE FROM boss_jobs")
            print("已清空现有数据")
            
            # 插入数据
            sql = """
            INSERT INTO boss_jobs (
                id, title, type, address, dist, educational, workExperience, 
                salary, salaryMonth, hrWork, hrName, companyTitle, companyNature, 
                companyStatus, companyPeople, companyAvatar, workTag, companyTags, 
                practice, detailUrl, companyUrl, createTime, user_id
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            
            # 批量插入数据
            data_tuples = []
            for index, row in df.iterrows():
                # 处理NaN值
                def handle_nan(value):
                    if pd.isna(value):
                        return None
                    return value
                
                data_tuple = (
                    handle_nan(row['id']), handle_nan(row['title']), handle_nan(row['type']), 
                    handle_nan(row['address']), handle_nan(row['dist']),
                    handle_nan(row['educational']), handle_nan(row['workExperience']), 
                    handle_nan(row['salary']), handle_nan(row['salaryMonth']),
                    handle_nan(row['hrWork']), handle_nan(row['hrName']), 
                    handle_nan(row['companyTitle']), handle_nan(row['companyNature']),
                    handle_nan(row['companyStatus']), handle_nan(row['companyPeople']), 
                    handle_nan(row['companyAvatar']), handle_nan(row['workTag']),
                    handle_nan(row['companyTags']), handle_nan(row['practice']), 
                    handle_nan(row['detailUrl']), handle_nan(row['companyUrl']),
                    handle_nan(row['createTime']), handle_nan(row['user_id'])
                )
                data_tuples.append(data_tuple)
            
            cursor.executemany(sql, data_tuples)
            
        connection.commit()
        print(f"成功导入{len(data_tuples)}条记录到数据库")
        
        # 验证导入结果
        with connection.cursor() as cursor:
            cursor.execute("SELECT COUNT(*) FROM boss_jobs")
            result = cursor.fetchone()
            if result:
                count = result[0]
                print(f"数据库中现有记录数: {count}")
            else:
                print("无法获取记录数")
            
    except Exception as e:
        print(f"导入数据时出错: {e}")
        connection.rollback()
    finally:
        connection.close()

def main():
    """主函数"""
    print("开始导入BOSS直聘数据到数据库...")
    
    # 创建数据表
    create_table()
    
    # 导入数据
    import_csv_to_db()
    
    print("数据导入完成")

if __name__ == "__main__":
    main()